from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import json
import asyncio
from app.config.config import API_KEY
from app.utils.logger import Logger
from app.utils.performance_monitor import start_timing, end_timing
from app.utils.async_tools import async_tool_and_summary, safe_send_json
from app.core.ai_core import AICore

# 会话管理：存储客户端ID到AI Core实例的映射
client_sessions = {}

# 会话超时时间（秒）
SESSION_TIMEOUT = 300

logger = Logger("WebSocketServer")
router = APIRouter()

@router.websocket("/ws")
async def websocket_chat_endpoint(websocket: WebSocket):
    # 验证 API Key
    api_key = websocket.headers.get("x-api-key") or websocket.query_params.get("api_key")
    if api_key != API_KEY:
        await websocket.close(code=4001, reason="Unauthorized")
        return

    await websocket.accept()
    logger.info("✅ WebSocket客户端已连接")

    # 从应用状态获取已初始化的AI Core实例
    app_ai_core = websocket.app.state.ai_core
    
    try:
        while True:
            data = await websocket.receive_text()
            logger.info(f"📩 收到客户端消息: {data}")

            try:
                message_data = json.loads(data)
            except json.JSONDecodeError:
                await safe_send_json(websocket, {
                    "type": "error",
                    "content": "消息格式错误，请发送JSON格式数据"
                })
                continue

            # 获取客户端ID，用于会话管理
            client_id = message_data.get("client_id", "default")
            query = (message_data.get("message") or message_data.get("query", "")).strip()
            stream = message_data.get("stream", True)
            enable_tools = message_data.get("enable_tools", True)

            if not query:
                await safe_send_json(websocket, {
                    "type": "error",
                    "content": "消息为空"
                })
                continue
            
            # 会话管理：获取或创建AI Core实例
            if client_id not in client_sessions:
                # 创建新的AI Core实例
                ai_core = AICore(
                    mcp_configs=app_ai_core.mcp_configs,
                    mcp_clients=app_ai_core.mcp_clients
                )
                # 复制工具列表
                ai_core.tools = app_ai_core.tools
                # 存储会话
                client_sessions[client_id] = {
                    "ai_core": ai_core,
                    "last_active": asyncio.get_event_loop().time()
                }
                logger.info(f"🧠 已为客户端 {client_id} 创建新的AI Core实例")
            else:
                # 更新会话最后活跃时间
                client_sessions[client_id]["last_active"] = asyncio.get_event_loop().time()
                ai_core = client_sessions[client_id]["ai_core"]
                logger.info(f"🧠 已复用客户端 {client_id} 的AI Core实例")
            
            # 清理超时会话
            current_time = asyncio.get_event_loop().time()
            expired_sessions = [cid for cid, session in client_sessions.items() 
                              if current_time - session["last_active"] > SESSION_TIMEOUT]
            for cid in expired_sessions:
                del client_sessions[cid]
                logger.info(f"🧹 已清理超时会话: {cid}")

            logger.info(f"🧠 处理用户查询: {query}")
            start_time = start_timing()

            # 立即发送处理开始的通知，让用户知道系统已收到请求
            await safe_send_json(websocket, {
                "type": "processing_start",
                "message": "正在处理您的请求..."
            })

            full_response = ""

            try:
                # 流式模式
                if stream:
                    chunk_count = 0
                    tool_called = False

                    async for msg in ai_core.process_query_stream(
                        query, enable_tools=enable_tools, websocket=websocket, use_memory=True
                    ):
                        if not msg:
                            continue

                        if isinstance(msg, (str, bytes)):
                            msg_str = msg.decode() if isinstance(msg, bytes) else msg
                            msg_str = msg_str.strip()
                            if not msg_str:
                                continue
                            # 特殊处理 "[完成]" 标记
                            if msg_str == "[完成]":
                                logger.info("准备发送 streaming_complete 消息")
                                await safe_send_json(websocket, {
                                    "type": "streaming_complete",
                                    "content": ""  # 在流式模式下，我们不收集完整响应，所以这里发送空字符串
                                })
                                logger.info("已发送 streaming_complete 消息")
                                logger.info(f"📤 流式发送完成，共 {chunk_count} 片段")
                                logger.info("🤖 AI最终结果: 流式处理完成")
                                break
                            data_json = {"type": "streaming", "content": msg_str}
                            full_response += msg_str
                        elif isinstance(msg, dict):
                            data_json = msg
                        else:
                            logger.warning(f"⚠️ 未知流式消息类型: {type(msg)} -> {msg}")
                            continue

                        # 检测到工具调用信息，直接跳过，让AI Core自己处理工具调用
                        # 在流式响应中，工具调用的arguments字段是逐步返回的，无法完整解析
                        if data_json.get("type") == "tool_call_info":
                            logger.info(f"📡 收到工具调用信息，跳过处理，由AI Core自行处理")
                            continue
                        
                        # 检测到普通工具调用（非流式）
                        if data_json.get("type") == "tool_call" and not tool_called:
                            tool_called = True
                            tool_name = data_json.get("tool_name")
                            tool_args = data_json.get("arguments") or data_json.get("tool_args", {})

                            # 立即通知前端工具调用
                            await safe_send_json(websocket, {
                                "type": "tool_call",
                                "tool_name": tool_name,
                                "arguments": tool_args,
                                "message": f"正在调用 {tool_name} 工具，请稍候..."
                            })
                            logger.info(f"🧰 触发工具调用: {tool_name}({tool_args})")

                            # 异步执行工具 + 后台生成总结
                            asyncio.create_task(
                                async_tool_and_summary(tool_name, tool_args, websocket, ai_core)
                            )
                            continue

                        # 常规AI输出
                        await safe_send_json(websocket, data_json)
                        chunk_count += 1

                    # 注意：完成消息已经在流式处理过程中发送，这里不再重复发送
                    # await safe_send_json(websocket, {
                    #     "type": "streaming_complete",
                    #     "content": full_response
                    # })

                    # logger.info(f"📤 流式发送完成，共 {chunk_count} 片段")
                    # logger.info(f"🤖 AI最终结果: {full_response[:120]}...")

                # 非流式模式
                else:
                    response = await ai_core.process_query(query, stream=False, enable_tools=enable_tools, update_memory=True)
                    await safe_send_json(websocket, {
                        "type": "response",
                        "content": response
                    })
                    logger.info(f"🤖 非流式AI结果: {response[:100]}...")
                    full_response = response
                    
                    # 发送完成标记
                    await safe_send_json(websocket, {
                        "type": "complete",
                        "content": ""
                    })

                # 在流式模式下，full_response 可能为空，因为我们不再收集完整响应
                # 在非流式模式下，full_response 应该包含完整的响应
                end_timing(start_time, query, full_response if stream else response, None)

            except Exception as e:
                logger.error(f"❌ 处理查询时出错: {e}")
                await safe_send_json(websocket, {
                    "type": "error",
                    "content": f"服务器处理错误: {str(e)}"
                })

    except WebSocketDisconnect:
        logger.info("🌐 客户端主动断开WebSocket连接")
    except Exception as e:
        logger.error(f"💥 WebSocket连接异常: {e}")

    finally:
        logger.info("🔴 WebSocket客户端已断开连接")
        # 每个客户端的AI Core实例会在连接断开后自动被垃圾回收

@router.websocket("/ws/asr")
async def websocket_asr_endpoint(websocket: WebSocket):
    """WebSocket 端点，用于接收音频数据并返回识别结果"""
    await websocket.accept()
    
    connection_id = str(id(websocket))
    # 从应用状态中获取活跃连接字典
    active_connections = websocket.app.state.active_connections if hasattr(websocket.app.state, 'active_connections') else {}
    active_connections[connection_id] = websocket
    # 更新应用状态中的活跃连接字典
    websocket.app.state.active_connections = active_connections
    
    logger.info(f"ASR模型已建立连接: {connection_id}")
    
    try:
        # 发送连接成功的消息
        await websocket.send_text(json.dumps({
            "type": "connection",
            "status": "connected",
            "message": "WebSocket 连接成功"
        }))
        
        # 从应用状态中获取ASR服务
        asr_service = websocket.app.state.asr_service
        if asr_service is None:
            await websocket.send_text(json.dumps({
                "type": "error",
                "text": "语音识别服务未初始化或初始化失败"
            }))
            return
        
        audio_buffer = bytearray()
        
        while True:
            # 接收客户端发送的数据
            data = await websocket.receive()
            
            if "bytes" in data:
                # 接收到二进制音频数据
                audio_chunk = data["bytes"]
                audio_buffer.extend(audio_chunk)
                
                # 每累积 30KB 数据进行一次识别（约1-2秒的音频）
                if len(audio_buffer) >= 30720:  # 30KB
                    # 进行语音识别
                    result = asr_service.recognize_audio(bytes(audio_buffer))
                    
                    logger.info(f"ASR识别结果: {result}")
                    
                    # 发送识别结果
                    await websocket.send_text(json.dumps({
                        "type": "result",
                        "text": result
                    }))
                    
                    # 清空缓冲区
                    audio_buffer.clear()
                    
            elif "text" in data:
                # 接收到文本消息
                message = data["text"]
                
                if message == "end":
                    # 处理最后一帧数据
                    if audio_buffer:
                        result = asr_service.recognize_audio(bytes(audio_buffer))
                        await websocket.send_text(json.dumps({
                            "type": "final",
                            "text": result
                        }))
                        audio_buffer.clear()
                        
                elif message == "ping":
                    # 心跳检测
                    await websocket.send_text(json.dumps({
                        "type": "pong"
                    }))
                    
    except WebSocketDisconnect:
        logger.info(f"WebSocket 连接断开: {connection_id}")
    except Exception as e:
        logger.error(f"WebSocket 处理错误: {e}")
        await websocket.send_text(json.dumps({
            "type": "error",
            "text": f"服务器错误: {str(e)}"
        }))
    finally:
        # 清理连接
        if connection_id in active_connections:
            del active_connections[connection_id]